#include <llvm/ExecutionEngine/Orc/LLJIT.h>
#include <llvm/ExecutionEngine/Orc/ThreadSafeModule.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/Support/Error.h>
#include <llvm/Support/InitLLVM.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/IR/Verifier.h>

using namespace llvm;
using namespace llvm::orc;

namespace {

Function *createAdd(Module &module, LLVMContext &ctx) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType = FunctionType::get(int32Ty, {int32Ty, int32Ty}, false);
  auto *fn = Function::Create(fnType, Function::InternalLinkage, "add", module);

  auto argsItr = fn->arg_begin();
  Argument *a = argsItr++;
  a->setName("a");
  Argument *b = argsItr++;
  b->setName("b");

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);
  Value *sum = builder.CreateNSWAdd(a, b, "sum");
  builder.CreateRet(sum);
  return fn;
}

Function *createClamp(Module &module, LLVMContext &ctx) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType = FunctionType::get(int32Ty, {int32Ty}, false);
  auto *fn =
      Function::Create(fnType, Function::InternalLinkage, "clamp_non_negative",
                       module);

  Argument &x = *fn->arg_begin();
  x.setName("x");

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);
  Value *isNegative = builder.CreateICmpSLT(&x, builder.getInt32(0), "is.neg");
  Value *sel =
      builder.CreateSelect(isNegative, builder.getInt32(0), &x, "clamped");
  builder.CreateRet(sel);
  return fn;
}

Function *createCompute(Module &module, LLVMContext &ctx, Function *addFn,
                        Function *clampFn) {
  IRBuilder<> builder(ctx);
  auto *int32Ty = builder.getInt32Ty();
  auto *fnType =
      FunctionType::get(int32Ty, {builder.getInt1Ty(), int32Ty}, false);
  auto *fn =
      Function::Create(fnType, Function::InternalLinkage, "compute", module);

  auto argsItr = fn->arg_begin();
  Argument *boost = argsItr++;
  boost->setName("boost");
  Argument *value = argsItr++;
  value->setName("value");

  BasicBlock *entry = BasicBlock::Create(ctx, "entry", fn);
  builder.SetInsertPoint(entry);
  Value *base = builder.CreateCall(clampFn, {value}, "base");
  Value *doubleBase =
      builder.CreateShl(base, builder.getInt32(1), "double.base");
  Value *sel =
      builder.CreateSelect(boost, doubleBase, base, "maybe.boosted");
  builder.CreateRet(sel);
  return fn;
}

void populateModule(Module &module, LLVMContext &ctx) {
  Function *addFn = createAdd(module, ctx);
  Function *clampFn = createClamp(module, ctx);
  createCompute(module, ctx, addFn, clampFn);

  // Expose compute as an externally callable function with C signature.
  Function *wrapper =
      Function::Create(FunctionType::get(Type::getInt32Ty(ctx),
                                         {Type::getInt1Ty(ctx),
                                          Type::getInt32Ty(ctx)},
                                         false),
                       Function::ExternalLinkage, "compute_entry", module);
  wrapper->setCallingConv(CallingConv::C);

  auto argsIt = wrapper->arg_begin();
  Argument *boost = argsIt++;
  boost->setName("boost");
  Argument *value = argsIt++;
  value->setName("value");

  IRBuilder<> builder(ctx);
  BasicBlock *entry = BasicBlock::Create(ctx, "entry", wrapper);
  builder.SetInsertPoint(entry);
  Function *computeFn = module.getFunction("compute");
  Value *result =
      builder.CreateCall(computeFn, {boost, value}, "result");
  builder.CreateRet(result);
}

Error runJIT() {
  auto jitExpected = LLJITBuilder().create();
  if (!jitExpected)
    return jitExpected.takeError();
  auto jit = std::move(*jitExpected);

  auto threadSafeCtx = std::make_unique<LLVMContext>();
  auto module = std::make_unique<Module>("jit-stage06", *threadSafeCtx);
  populateModule(*module, *threadSafeCtx);

  if (verifyModule(*module, &errs()))
    return createStringError(inconvertibleErrorCode(),
                             "Module verification failed");

  ThreadSafeModule tsm(std::move(module), std::move(threadSafeCtx));
  if (auto err = jit->addIRModule(std::move(tsm)))
    return err;

  auto symExpected = jit->lookup("compute_entry");
  if (!symExpected)
    return symExpected.takeError();

  auto computePtr =
      symExpected->toPtr<int32_t (*)(bool, int32_t)>();

  int32_t boosted = computePtr(true, -3);
  int32_t plain = computePtr(false, 5);

  outs() << "compute(true, -3) = " << boosted << "\n";
  outs() << "compute(false, 5) = " << plain << "\n";
  return Error::success();
}

} // namespace

int main(int argc, char **argv) {
  InitLLVM initLLVM(argc, argv);
  InitializeNativeTarget();
  InitializeNativeTargetAsmPrinter();

  if (auto err = runJIT()) {
    errs() << "JIT execution failed: " << toString(std::move(err)) << "\n";
    return 1;
  }
  return 0;
}
